package org.eoti.ai.bayesian.example1;

import org.eoti.ai.bayesian.BayesianRule;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.concurrent.ConcurrentHashMap;

public class WordRule
extends BayesianRule<String>
{
    private ConcurrentHashMap<String,Word> words;
    private InterestComparator interestComparator;
    private int badTotal, goodTotal;

    public WordRule()
    {
        words = new ConcurrentHashMap<String,Word>();
        interestComparator = new InterestComparator();
        badTotal = 0;
        goodTotal = 0;
    }
    
    @Override
    public void train(String data, boolean isGoodData) {
        int count = (isGoodData ? goodTotal : badTotal);
        for(String text : data.split("\\W"))
        {
            text = text.trim().toLowerCase();
            if(text.length() < 3) continue;
            if(words.containsKey(text))
            {
                Word word = words.get(text);
                if(isGoodData) word.incrementGood();
                else word.incrementBad();
            }else{
                Word word = new Word(text);
                if(isGoodData) word.incrementGood();
                else word.incrementBad();
                words.put(text, word);
            }
            count++;
        }

        for(Word word : words.values())
        {
            if(isGoodData) word.updateGoodRate(count);
            else word.updateBadRate(count);
            word.updateSpamProbability();
        }

        if(isGoodData) goodTotal = count;
        else badTotal = count;
    }

    @Override
    protected double probability(String data) {
        ArrayList<Word> interesting = new ArrayList<Word>();
        
        for(String text : data.split("\\W"))
        {
            text = text.trim().toLowerCase();
            if(text.length() < 3) continue;
            Word word = words.get(text);
            if(word == null)
            {
                word = new Word(text);
                word.setSpamProbability(0.4);
            }
            interesting.add(word);
        }
//        Collections.sort(interesting, interestComparator);
        double pos = 1.0;
        double neg = 1.0;
//        System.out.format("Most interesting:\n");
        for(Word word : interesting)
//        for(int i=0; i<Math.min(15,interesting.size()); i++)
        {
//            Word word = interesting.get(i);
            double prob = word.getSpamProbability();
//            System.out.format("\t%s: %s\n", word.getText(), prob);
            pos *= prob;
            neg *= (1.0-prob);
        }
        double probability = pos / (pos+neg);
        if(probability < MIN_PROBABILITY)
            probability = MIN_PROBABILITY;

        if(probability > MAX_PROBABILITY)
            probability = MAX_PROBABILITY;

        return probability;
    }

    static class InterestComparator
    implements Comparator<Word>
    {
        @Override
        public int compare(Word w1, Word w2) {
            // we actually want it reversed
//            return ((Double)w1.getSpamProbability()).compareTo((Double)w2.getSpamProbability());
            return ((Double)w2.getSpamProbability()).compareTo((Double)w1.getSpamProbability());
        }
    }

    // derived from http://www.shiffman.net/teaching/a2z/bayesian/
    static class Word
    {
        private String text;
        private int countGood, countBad;
        private double rateGood, rateBad;
        private double spamProbability;

        public Word(String text)
        {
            this.text = text;
            countGood = 0;
            countBad = 0;
            rateGood = 0.0;
            rateBad = 0.0;
            spamProbability = 0.0;
        }

        public void incrementBad(){countBad++;}
        public void incrementGood(){countGood++;}
        public void updateBadRate(int totalBad)
        {
            if(totalBad <= 0) return;
            rateBad = (double)countBad / (double)totalBad;
//            System.out.format("%d / %d = %s BAD\n", countBad, totalBad, rateBad);
        }
        public void updateGoodRate(int totalGood)
        {
            if(totalGood <= 0) return;
            // multiply by 2 to help fight false positives per Graham and Daniel
            rateGood = (2.0 * (double)countGood) / (double)totalGood;
//            System.out.format("(2.0 * %d) / %d = %s GOOD\n", countGood, totalGood, rateGood);
        }
        public void updateSpamProbability()
        {
            if( (rateGood + rateBad) > 0.0 )
                spamProbability = rateBad / (rateBad + rateGood);

            if(spamProbability < MIN_PROBABILITY)
                spamProbability = MIN_PROBABILITY;

            if(spamProbability > MAX_PROBABILITY)
                spamProbability = MAX_PROBABILITY;

//            System.out.format("%s / (%s + %s) = %s\n", rateBad, rateBad, rateGood, spamProbability);
        }

        // how far off from 50/50
        public double howInteresting()
        {
            return Math.abs(0.5 - spamProbability);
        }

        public void setSpamProbability(double probability)
        {
            spamProbability = probability;
        }

        public String getText(){return text;}
        public int getGoodCount(){return countGood;}
        public int getBadCount(){return countBad;}
        public double getGoodRate(){return rateGood;}
        public double getBadRate(){return rateBad;}
        public double getSpamProbability(){return spamProbability;}
    }
}
